package top.szhome.framework.core.tools.util;

import top.szhome.framework.core.tools.datastruct.Tree;

import javax.validation.constraints.NotNull;
import java.util.LinkedList;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;

/**
 * 树形工具
 *
 * @author 唐加利
 * @date 2023/8/4
 **/
public class TreeUtil {

    /**
     * 列表转树形列表
     *
     * @param nodes 节点列表
     * @param <T>   节点ID类型
     * @param <N>   节点类型
     * @return  返回所有已设置字节点的根节点
     */
    public static <T, N extends Tree<N, T>> List<N> tree(List<N> nodes) {
        // 缓存全部节点
        Map<T, N> map = nodes.stream().collect(Collectors.toMap(Tree::getId, node -> node));

        List<N> result = new LinkedList<>();
        for(N node : nodes){
            T pId = node.getParentId();
            N parent = map.get(pId);

            // 没有父节点，作为根节点返回
            if(pId == null || parent == null){
                result.add(node);
                continue;
            }

            // 初始化字节点
            List<N> children = parent.getChildren();
            if(ObjectUtil.isEmpty(children)){
                children = new LinkedList<>();
                parent.setChildren(children);
            }

            // 添加到对应父节点的字节点列表
            children.add(node);
        }

        return result;
    }

    /**
     * 获取树形节点列表中指定节点的子节点
     *
     * @param nodeId    指定节点
     * @param nodes     树形节点列表
     * @param <T>       节点ID类型
     * @param <N>       节点类型
     * @return  返回树形节点列表中，给定ID对应节点的子节点
     */
    public static <T, N extends Tree<N, T>> List<N> getChildren(@NotNull Long nodeId, List<N> nodes) {
        if(nodeId == null || ObjectUtil.isEmpty(nodes)){
            return null;
        }

        for(N node : nodes){
            // 匹配到指定节点，返回其字节点
            if(ObjectUtil.equals(nodeId, node.getId())){
                return node.getChildren();
            }

            List<N> children = getChildren(nodeId, node.getChildren());
            if(children != null){
                return children;
            }
        }

        return null;
    }

}
